clc;
clear;
close all;

addpath('autogen\')
%%%%%%%%%%%%%% 重力项标定 %%%%%%%%%%%%%%

% CAD评估的结果
m1=0.47; m2=0.29; m3=0.24; m4=0.23; m5=0.08;
r1_x=0.00; r1_y=0.00; r1_z=-0.01;
r2_x=0.01; r2_y=-0.07; r2_z=0.0;
r3_x=0.00; r3_y=0.00; r3_z=-0.08;
r4_x=0.00; r4_y=-0.06; r4_z=0.00;
r5_x=0.00; r5_y=0.00; r5_z=-0.02;

lb = [0, 0, 0, 0, 0,...
    r1_x-0.0, r1_y-0.0, r1_z-0.1,...
    r2_x-0.01, r2_y-0.1, r2_z-0.0,...
    r3_x-0.0, r3_y-0.0, r3_z-0.1,...
    r4_x-0.0, r4_y-0.1, r4_z-0.0,...
    r5_x-0.0, r5_y-0.0, r5_z-0.1];

ub = [m1+1.0, m2+1.0, m3+1.0, m4+1.0, m5+1.0,...
    r1_x+0.0, r1_y+0.0, r1_z+0.01,...
    r2_x+0.1, r2_y+0.07, r2_z+0.0,...
    r3_x+0.0, r3_y+0.0, r3_z+0.08,...
    r4_x+0.0, r4_y+0.06, r4_z+0.0,...
    r5_x+0.0, r5_y+0.0, r5_z+0.02];

m5_r5 = m5*[r5_x, r5_y, r5_z, 1]';
m4_r4 = m4*[r4_x, r4_y, r4_z, 1]';
m3_r3 = m3*[r3_x, r3_y, r3_z, 1]';
m2_r2 = m2*[r2_x, r2_y, r2_z, 1]';
m1_r1 = m1*[r1_x, r1_y, r1_z, 1]';

% This Params has been deployed.
% m1_r1 = [0;0;-0.00470000000000000;0.470000000000000];
% m2_r2 = [1.83003286952368e-05;-0.0910502644964221;0;0.636003391887722];
% m3_r3 = [0;0;-0.00316201862885600;0.230305754293214];
% m4_r4 = [0;-0.0129295831211659;0;0.286035027497600];
% m5_r5 = [0;0;-0.00460183240367810;0.141954767292057];

g = [0 0 -9.8 0]';

% 读取数据
data = load('data_l3/r_arm.txt');
% data = data(1:298, :);

alpha = [pi/2, -pi/2, pi/2, pi/2, pi/2]';

qs = data(:, 1:5);

tor = data(:, 11:15);
tor(1:398, 2) = tor(1:398, 2)*-1.0;
% load('tau_3s_est.mat');
% tor(:, 3) = tau_3s_est;
tor(:, 5) = tor(:, 5)*0.0;


K = [];
Tau = [];

tau_1s = [];tau_2s = [];tau_3s = [];tau_4s = [];tau_5s = [];
tau_1s_est = [];tau_2s_est = [];tau_3s_est = [];tau_4s_est = [];tau_5s_est = [];
for i=1:1:length(qs)
    q = qs(i, :)'; 
    
    if qs(2) > -3.14
        qs(i, 2) = -(q(2)+3.14);
    else
        qs(i, 2) = abs(q(2)+3.14);
    end 
    
    u11 = J_T01_q1(q, alpha);u21 = J_T02_q1(q, alpha);u31 = J_T03_q1(q, alpha);u41 = J_T04_q1(q, alpha);u51 = J_T05_q1(q, alpha);
    u22 = J_T02_q2(q, alpha);u32 = J_T03_q2(q, alpha);u42 = J_T04_q2(q, alpha);u52 = J_T05_q2(q, alpha);
    u33 = J_T03_q3(q, alpha);u43 = J_T04_q3(q, alpha);u53 = J_T05_q3(q, alpha);
    u44 = J_T04_q4(q, alpha);u54 = J_T05_q4(q, alpha);
    u55 = J_T05_q5(q, alpha);
    
    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4-g'*u51*m5_r5;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4-g'*u52*m5_r5;
    tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4-g'*u53*m5_r5;
    tau_4 = -g'*u44*m4_r4-g'*u54*m5_r5;
    tau_5 = -g'*u55*m5_r5;

    O = [0, 0, 0, 0];
    Kj = [
        g'*u11, g'*u21, g'*u31, g'*u41, g'*u51;
        O, g'*u22, g'*u32, g'*u42, g'*u52;
        O, O, g'*u33, g'*u43, g'*u53;
        O, O, O, g'*u44, g'*u54;
        O, O, O, O, g'*u55
        ];

    K = [K; Kj];
    
    tau_5s = [tau_5s; tor(i, 5)];
    tau_4s = [tau_4s; tor(i, 4)];
    tau_3s = [tau_3s; tor(i, 3)];
    tau_2s = [tau_2s; tor(i, 2)];
    tau_1s = [tau_1s; tor(i, 1)];
    
    Tau = [Tau; 
        tor(i, 1)'; 
        tor(i, 2)'; 
        tor(i, 3)'; 
        tor(i, 4)'; 
        tor(i, 5)'
        ];
    
    tau_5s_est = [tau_5s_est; tau_5];
    tau_4s_est = [tau_4s_est; tau_4];
    tau_3s_est = [tau_3s_est; tau_3];
    tau_2s_est = [tau_2s_est; tau_2];
    tau_1s_est = [tau_1s_est; tau_1];
end

% figure(1);
% plot(tau_1s, "r.-");
% hold on;
% plot(tau_1s_est, "g.-");
% legend('real', 'est');   
% figure(2);
% plot(tau_2s, "r.-");
% hold on;
% plot(tau_2s_est, "g.-");
% legend('real', 'est');   
% figure(3);
% plot(tau_3s, "r.-");
% hold on;
% plot(tau_3s_est, "g.-");
% legend('real', 'est');   
% figure(4);
% plot(tau_4s, "r.-");
% hold on;
% plot(tau_4s_est, "g.-");
% legend('real', 'est');   
% figure(5);
% plot(tau_5s, "r.-");
% hold on;
% plot(tau_5s_est, "g.-");
% legend('real', 'est');   


% 非线性最小二乘估计
x_init = [
    m1, m2, m3, m4, m5,...
    r1_x, r1_y, r1_z, ...
    r2_x, r2_y, r2_z, ...
    r3_x, r3_y, r3_z, ...
    r4_x, r4_y, r4_z, ...
    r5_x, r5_y, r5_z];

qs_input = qs;
Taus = tau_estimator(x_init, qs);

% x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau);
x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau, lb, ub);

% m1_r1 = [x(1)*x(6); x(1)*x(7); x(1)*x(8); x(1)]; %[0;0;-0.00470000000000000;0.470000000000000]
% m2_r2 = [x(2)*x(9); x(2)*x(10); x(2)*x(11); x(2)]; %[0.00433456455771394;-0.0361043073452328;0;0.427188331792686]
% m3_r3 = [x(3)*x(12); x(3)*x(13); x(3)*x(14); x(3)]; %[0;0;-0.0262620924814664;0.384060775592276]
% m4_r4 = [x(4)*x(15); x(4)*x(16); x(4)*x(17); x(4)]; %[0;-0.00602882538693720;0;0.356727708453166]
% m5_r5 = [x(5)*x(18); x(5)*x(19); x(5)*x(20); x(5)]; %[0;0;-0.00528331435735030;0.192142299786217]
    
tau_1s_real = [];tau_2s_real = [];tau_3s_real = [];tau_4s_real = [];tau_5s_real = [];
tau_1s_est = [];tau_2s_est = [];tau_3s_est = [];tau_4s_est = [];tau_5s_est = [];
for i=1:1:length(qs_input)
    q = qs_input(i, :)'; 

    u11 = J_T01_q1(q, alpha);u21 = J_T02_q1(q, alpha);u31 = J_T03_q1(q, alpha);u41 = J_T04_q1(q, alpha);u51 = J_T05_q1(q, alpha);
    u22 = J_T02_q2(q, alpha);u32 = J_T03_q2(q, alpha);u42 = J_T04_q2(q, alpha);u52 = J_T05_q2(q, alpha);
    u33 = J_T03_q3(q, alpha);u43 = J_T04_q3(q, alpha);u53 = J_T05_q3(q, alpha);
    u44 = J_T04_q4(q, alpha);u54 = J_T05_q4(q, alpha);
    u55 = J_T05_q5(q, alpha);

    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4-g'*u51*m5_r5;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4-g'*u52*m5_r5;
    tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4-g'*u53*m5_r5;
    tau_4 = -g'*u44*m4_r4-g'*u54*m5_r5;
    tau_5 = -g'*u55*m5_r5;
    
    tau_5s_real = [tau_5s_real; tor(i, 5)];
    tau_4s_real = [tau_4s_real; tor(i, 4)];
    tau_3s_real = [tau_3s_real; tor(i, 3)];
    tau_2s_real = [tau_2s_real; tor(i, 2)];
    tau_1s_real = [tau_1s_real; tor(i, 1)];
    
    tau_5s_est = [tau_5s_est; tau_5];
    tau_4s_est = [tau_4s_est; tau_4];
    tau_3s_est = [tau_3s_est; tau_3*0.5];
    tau_2s_est = [tau_2s_est; tau_2*1.6];
    tau_1s_est = [tau_1s_est; tau_1*1.6];
end

% plot
figure(1);
plot(tau_1s_real, "r.-");
hold on;
plot(tau_1s_est, "g.-");
figure(2);
plot(tau_2s_real, "r.-");
hold on;
plot(tau_2s_est, "g.-");
figure(3);
plot(tau_3s_real, "r.-");
hold on;
plot(tau_3s_est, "g.-");
figure(4);
plot(tau_4s_real, "r.-");
hold on;
plot(tau_4s_est, "g.-");
figure(5);
plot(tau_5s_real, "r.-");
hold on;
plot(tau_5s_est, "g.-");
legend('real', 'est');   


function taus = tau_estimator(x, qs_input)
    alpha = [pi/2, -pi/2, pi/2, pi/2, pi/2]';
    g = [0 0 -9.8 0]';
    
    m1_r1 = [x(1)*x(6); x(1)*x(7); x(1)*x(8); x(1)];
    m2_r2 = [x(2)*x(9); x(2)*x(10); x(2)*x(11); x(2)];
    m3_r3 = [x(3)*x(12); x(3)*x(13); x(3)*x(14); x(3)];
    m4_r4 = [x(4)*x(15); x(4)*x(16); x(4)*x(17); x(4)];
    m5_r5 = [x(5)*x(18); x(5)*x(19); x(5)*x(20); x(5)];
    
    taus = [];
    for i=1:1:length(qs_input)
        q = qs_input(i, :)'; 

        u11 = J_T01_q1(q, alpha);u21 = J_T02_q1(q, alpha);u31 = J_T03_q1(q, alpha);u41 = J_T04_q1(q, alpha);u51 = J_T05_q1(q, alpha);
        u22 = J_T02_q2(q, alpha);u32 = J_T03_q2(q, alpha);u42 = J_T04_q2(q, alpha);u52 = J_T05_q2(q, alpha);
        u33 = J_T03_q3(q, alpha);u43 = J_T04_q3(q, alpha);u53 = J_T05_q3(q, alpha);
        u44 = J_T04_q4(q, alpha);u54 = J_T05_q4(q, alpha);
        u55 = J_T05_q5(q, alpha);
        
        tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4-g'*u51*m5_r5;
        tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4-g'*u52*m5_r5;
        tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4-g'*u53*m5_r5;
        tau_4 = -g'*u44*m4_r4-g'*u54*m5_r5;
        tau_5 = -g'*u55*m5_r5;
        
        taus = [
            taus; 
            tau_1; 
            tau_2; 
            tau_3; 
            tau_4;
            tau_5
            ];
    end
end
